package dmir.lda.adapters;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Properties;
import java.util.Set;

/**
 * 
 * Implementation of an LdaAdapter for GibbsLDA++ (http://gibbslda.sourceforge.net)
 * 
 * @author ianastacio
 *
 */
public class GibbsLda implements LdaAdapter{

	private static final int THETA_FILE_LINE_LENGTH = 901;
	private static final int PHI_FILE_TOKEN_LENGTH = 8;
	private static final int PHI_FILE_LINE_LENGTH = 1480573;
	private static final int PHI_FILE_SPACES_LENGTH = 1;
	
    private final RandomAccessFile phiRaf;
    private final RandomAccessFile thetaRaf;
    private final RandomAccessFile tassignRaf;
	
	private final int numTokens;
	private final int numDocs;
	private final int numTopics;
	
	private Map<Integer, Long> tassignDocOffsets;
	
	/**
	 * Creates an GibbsLda++ adapter based on the files present in the given directory.
	 * 
	 * @param dir The directory containing the GibbsLda++ files
	 * @param modelName The basename for the LDA model produced by GibbsLDA++ (eg. model-final).
	 * @throws Exception 
	 */
	public GibbsLda(File dir, String modelName) throws Exception {
		
		///////////////////
		
		Properties props = new Properties();
		FileInputStream in = new FileInputStream(new File(dir, modelName + ".others"));
		props.load(in);
		in.close();
		
		numTokens = new Integer( props.getProperty("nwords") );
		numDocs = new Integer( props.getProperty("ndocs") );
		numTopics = new Integer( props.getProperty("ntopics") );
		
		///////////////////
		
		phiRaf = new RandomAccessFile(new File(dir, modelName + ".phi"), "r");
		thetaRaf = new RandomAccessFile(new File(dir, modelName + ".theta"), "r");
		tassignRaf = new RandomAccessFile(new File(dir, modelName + ".tassign"), "r");
	}

	@Override
	public int numTokens() {
		return numTokens;
	}

	@Override
	public int numDocuments() {
		return numDocs;
	}

	@Override
	public int numTopics() {
		return numTopics;
	}

	@Override
	public Set<Integer> docTokens(int doc) {
		
		Set<Integer> tokens = new HashSet<Integer>();
		
		try {
			if (tassignDocOffsets == null) computeTassignOffsets();

			tassignRaf.seek( tassignDocOffsets.get(doc) );
			String line = tassignRaf.readLine();
			String[] tokenStrArray = line.split(" ");
			
			for (int i = 0; i < tokenStrArray.length; i++) 
			{
				tokens.add( new Integer( tokenStrArray[i].split(":")[0] ) );
			}
			
		} catch (IOException e) {
			e.printStackTrace();
		}

		return tokens;
	}

	
	private void computeTassignOffsets() throws IOException {
		
		tassignDocOffsets = new HashMap<Integer, Long>(numDocs);
		
		long counter = 0;
		int doc = 0;
		String line = "";
		while ((line = tassignRaf.readLine()) != null) 
		{
			tassignDocOffsets.put(doc, counter);
			
			doc++;
			counter += line.length() + 1;
		}
	}

	@Override
	public double[] tokenProbabilities(int token) {
		
		double[] probabilities = new double[numTopics];
		
		for (int topic = 0; topic < numTopics; topic++) 
		{
			try 
			{
				phiRaf.seek(
						(long)topic * PHI_FILE_LINE_LENGTH 
						+ (long)token * PHI_FILE_TOKEN_LENGTH 
						+ (long)token * PHI_FILE_SPACES_LENGTH);

				probabilities[topic] = phiRaf.readDouble();
				
			} catch (IOException e) {
				e.printStackTrace();
			}
		}

		return probabilities;
	}

	
	@Override
	public double[] docProbabilities(int doc) {
		
		double[] probabilities = new double[numTopics];
		
		try 
		{
			thetaRaf.seek(THETA_FILE_LINE_LENGTH * (long)doc);

			byte[] b = new byte[THETA_FILE_LINE_LENGTH];
			thetaRaf.read(b, 0, b.length);
			String line = new String(b);

			String[] probStrings = line.trim().split(" ");
			for (int i = 0; i < probStrings.length; i++) 
			{
				probabilities[i] = Double.parseDouble(probStrings[i]);
			}
		} catch (IOException e) {
			e.printStackTrace();
		}

		return probabilities;
	}

}
